// mcts.h
#ifndef MCTS_H
#define MCTS_H

#include "base.h"
#include <memory>
#include <vector>
#include <unordered_map>
#include <random>
#include <cmath>
#include <iostream>

/**
 * @brief Forward declaration of GameState.
 */
class GameState;

/**
 * @brief Class representing a node in the MCTS tree.
 */
class MCTSNode : public std::enable_shared_from_this<MCTSNode> {
public:
    GameState game_state;                              ///< Current game state at this node
    std::weak_ptr<MCTSNode> parent;                    ///< Parent node in the tree
    Move move;                                         ///< Move that led to this node
    std::unordered_map<Player, int> win_counts;       ///< Win counts for each player
    int num_rollouts;                                  ///< Number of simulations run through this node
    std::vector<std::shared_ptr<MCTSNode>> children;   ///< Child nodes
    std::vector<Move> unvisited_moves;                 ///< Moves that have not been expanded yet

    /**
     * @brief Constructor for MCTSNode.
     * 
     * @param game_state_ Current game state.
     * @param parent_ Parent node.
     * @param move_ Move that led to this node.
     */
    MCTSNode(const GameState& game_state_, std::shared_ptr<MCTSNode> parent_ = nullptr, const Move& move_ = Move())
        : game_state(game_state_), parent(parent_), move(move_), 
          win_counts{ {Player::Black, 0}, {Player::White, 0} }, num_rollouts(0), 
          children(), unvisited_moves(game_state.legal_moves()) {}

    /**
     * @brief Adds a random child node by expanding an unvisited move.
     * 
     * @return std::shared_ptr<MCTSNode> Newly added child node.
     */
    std::shared_ptr<MCTSNode> add_random_child();

    /**
     * @brief Records a win for the given player.
     * 
     * @param winner Player who won the simulation.
     */
    void record_win(Player winner);

    /**
     * @brief Checks if the node can be expanded further.
     * 
     * @return true If there are unvisited moves.
     * @return false Otherwise.
     */
    bool can_add_child() const;

    /**
     * @brief Checks if the game state at this node is terminal.
     * 
     * @return true If the game is over.
     * @return false Otherwise.
     */
    bool is_terminal() const;

    /**
     * @brief Calculates the winning fraction for a given player.
     * 
     * @param player Player for whom to calculate the winning fraction.
     * @return double Winning fraction.
     */
    double winning_frac(Player player) const;
};

/**
 * @brief Calculates the UCT (Upper Confidence Bound applied to Trees) score.
 * 
 * @param parent_rollouts Total number of simulations from the parent node.
 * @param child_rollouts Number of simulations from the child node.
 * @param winning_pct Winning percentage of the child node.
 * @param temperature Exploration parameter.
 * @return double UCT score.
 */
inline double uct_score(int parent_rollouts, int child_rollouts, double winning_pct, double temperature) {
    if (child_rollouts == 0) {
        return std::numeric_limits<double>::infinity();
    }
    double exploration = std::sqrt(std::log(static_cast<double>(parent_rollouts)) / static_cast<double>(child_rollouts));
    return winning_pct + temperature * exploration;
}

/**
 * @brief Class implementing the MCTS Agent.
 */
class MCTSAgent : public Agent {
public:
    /**
     * @brief Constructor for MCTSAgent.
     * 
     * @param num_rounds_ Number of MCTS iterations.
     * @param temperature_ Exploration parameter for UCT.
     */
    MCTSAgent(int num_rounds_, double temperature_) 
        : Agent(num_rounds_, temperature_), rng(std::random_device{}()) {}

    /**
     * @brief Selects the best move based on MCTS.
     * 
     * @param game_state Current state of the game.
     * @return Move Selected move.
     */
    virtual Move select_move(const GameState& game_state) override;

private:
    std::mt19937 rng; ///< Random number generator.

    /**
     * @brief Selects the best child node based on UCT score.
     * 
     * @param node Current node.
     * @return std::shared_ptr<MCTSNode> Selected child node.
     */
    std::shared_ptr<MCTSNode> select_child(std::shared_ptr<MCTSNode> node);

    /**
     * @brief Simulates a random game from the given game state.
     * 
     * @param game_state Game state to start the simulation.
     * @return Player Winner of the simulated game.
     */
    Player simulate_random_game(const GameState& game_state);
};

#endif // MCTS_H
